from itertools import chain
import platform,random,string
import numpy as np
from sumolib.net import readNet,Net
from sumolib.net.edge import Edge
from feasytools import PQueue
from .cslist import *
from .ev import *
from .params import WINDOWS_VISSUALIZE

if platform.system()=="Linux":
    import libsumo as traci
else: # Windows & Mac
    if WINDOWS_VISSUALIZE:
        import traci
    else:
        import libsumo as traci

from traci._simulation import Stage
TC = traci.constants
TWeights = tuple[float,float,float]

_letters = string.ascii_letters + string.digits
def _random_string(length:int):
    return ''.join(random.choice(_letters) for _ in range(length))

class TrafficInst:    
    @staticmethod
    def __find_route(e1:str,e2:str)->Stage:
        return traci.simulation.findRoute(e1,e2,routingMode=TC.ROUTING_MODE_AGGREGATED)

    def __find_route_static(self,e1:str,e2:str)->Stage:
        return self.__shortest_paths[f"{e1}|{e2}"]
    
    def __init__(self, road_net_file:str, evs:Union[str,EVDict], start_time:int, end_time:int, seed:int = 0,
        cslist:Union[None,CSList,str]=None, pklist:Union[None,CSList,str]=None, ostream=None):
        '''
        路网模型初始化
            road_net_file: SUMO路网配置文件
            evs: 车辆信息与行程, 可以是文件名, 也可以是EVDict
            seed: 随机化种子
            n_threads: 并行处理线程数
            cslist: 快充站列表, None表示所有CS开头的道路都是充电站, 字符串表示从该文件读入
            pklist: 支持V2G的停车点列表, None表示所有非CS开头的道路都支持V2G, 字符串表示从该文件读入
            ostream: 指示打印的输出流, None表示stdout
        '''
        random.seed(seed)
        self.__seed=seed
        self.__gui=None
        self.__ctime:int = start_time
        self.__stime:int = start_time
        self.__etime:int = end_time
        self.__ostream = ostream
        # 读取路网
        self.__pendSet:set = set()
        self.__driveSet:set = set()
        self.__rnet:Net = readNet(road_net_file)
        self.__edges:list[Edge] = self.__rnet.getEdges()
        self.__names:list[str] = [e.getID() for e in self.__edges]
        self.__names_cs:list[str] = []
        self.__names_ncs:list[str] = []
        
        for e in self.__names:
            if not e.startswith("CS"):
                self.__names_ncs.append(e)
            elif cslist is None:
                if not e.endswith("rev"): self.__names_cs.append(e)
        
        # 加载静态最短路径
        self.__shortest_paths:dict[str,Stage]={}
        for e1 in self.__edges:
            for e2 in self.__edges:
                if e1.getID()==e2.getID(): continue
                route,cost=self.__rnet.getShortestPath(e1,e2)
                route:list[Edge]
                self.__shortest_paths[f"{e1.getID()}|{e2.getID()}"]=Stage(
                    edges=[e.getID() for e in route],
                    length=sum(e.getLength() for e in route),
                    travelTime=sum(e.getLength()/e.getSpeed() for e in route)
                )

        # 加载车辆
        self._fQ = PQueue() # 故障队列
        self._que = PQueue() # 发车队列
        self._VEHs = evs if isinstance(evs,EVDict) else EVDict(evs)

        # 加载充电站
        if isinstance(cslist,str): cslist = CSList(self._VEHs,filePath=cslist,csType=FCS)
        if cslist is None:            
            self._cs:CSList[FCS] = CSList(self._VEHs,csList=[FCS(e,DEFAULT_CS_SLOTS,DEFAULT_BUS_NAME,DEFAULT_BUY_PRICE) for e in self.__names_cs])
            print(f"快充站设置为缺省值, 共{len(self._cs)}座.",file=self.__ostream)
        elif isinstance(cslist, CSList):
            self._cs:CSList[FCS] = cslist
            self.__names_cs = cslist.get_CS_names()
        else:
            raise ValueError("cslist类型错误")
        
        if isinstance(pklist,str): pklist = CSList(self._VEHs,filePath=pklist,csType=SCS)
        if pklist is None:
            self._pk:CSList[SCS] = CSList(self._VEHs,csList=[SCS(e,DEFAULT_PK_SLOTS,DEFAULT_BUS_NAME,DEFAULT_BUY_PRICE,DEFAULT_SELL_PRICE) for e in self.__names_ncs])
            print(f"路边慢充设置为缺省值, 共{len(self._pk)}座.",file=self.__ostream)
        elif isinstance(pklist, CSList):
            self._pk:CSList[SCS] = pklist
        else:
            raise ValueError("pklist类型错误")
        
        # 将车辆加载到充电站并准备发车
        for veh in self._VEHs.values():
            self._que.push(veh.trip.depart_time,veh.ID)
            # 有20%的概率添加到可充电停车点
            if veh.SOC<veh.ksc or random.random()<=0.2:
                self.__start_charging_SCS(veh)

    @property
    def start_time(self):
        '''仿真开始时间'''
        return self.__stime
    
    @property
    def end_time(self):
        '''仿真结束时间'''
        return self.__etime
    
    @property
    def current_time(self):
        '''当前时间'''
        return self.__ctime
    
    @property
    def CSList(self):
        '''快充站列表'''
        return self._cs
    
    @property
    def PKList(self):
        '''快充站列表'''
        return self._pk
    
    @property
    def vehicles(self)->EVDict:
        '''车辆字典, key为车辆ID, value为EV实例'''
        return self._VEHs
    
    def __add_veh(self,veh_id:str,route:list[str]):
        self._VEHs[veh_id].clear_odometer()
        rou_id=_random_string(16)
        traci.route.add(rou_id, route)
        traci.vehicle.add(veh_id,rou_id)
    
    @property
    def edges(self)->list[Edge]:
        '''获取所有道路'''
        return self.__edges
    
    @property
    def trips_iterator(self):
        '''获取所有行程的迭代器'''
        return chain(*(x.trips for x in self._VEHs.values()))
    
    def get_edge_names(self, contain_CS = False)->list[str]:
        '''
        获取所有道路的名称
            contain_CS: 是否包括充电站道路
        '''
        return self.__names if contain_CS else self.__names_ncs
    
    def __sel_best_CS(self, veh:EV, omega:float, current_edge:Optional[str]=None)->tuple[list[str],TWeights]:
        '''
        根据汽车当前所在的edge选择离它最近的可用充电站, 返回路径和平均权重
            veh: 车辆实例
            omega: 权重
            current_edge: 当前道路, 如果为None则自动获取
        返回:
            路径(list[str]), 权重(tuple[float,float,float])
            如果没有找到充电站，返回[],(-1,-1,-1)
        '''
        to_charge=veh.charge_target-veh.battery
        c_edge = traci.vehicle.getRoadID(veh.ID) if current_edge is None else current_edge

        # 距离检查
        cs_names:list[str] = []; veh_cnt:list[int] = []; slots:list[int] = []
        prices:list[float] = []; stages:list[Stage] = []
        for cs in self._cs:
            stage=self.__find_route(c_edge,cs.name)
            if cs.is_online(self.__ctime) and veh.is_batt_enough(stage.length):
                cs_names.append(cs.name)
                veh_cnt.append(cs.veh_count())
                slots.append(cs.slots)
                prices.append(cs.pbuy(self.__ctime))
                stages.append(stage)

        if len(cs_names)==0:
            return [], (-1,-1,-1)
        
        t_drive = np.array([t.travelTime for t in stages]) / 60 # 将行驶时间转化为分钟
        t_wait = np.array([max(t-lim,0) for t,lim in zip(veh_cnt,slots)]) * 30 # 排队时间: 分钟
        #+ to_charge/(veh._efc_rate*60) # 充电时间：假定30min/辆
        
        # 总权重
        weight = np.sum([
            omega * (t_drive + t_wait),     # 驾驶时间和排队时间权重
            to_charge * np.array(prices)    # 电价权重
        ], axis=0).tolist()

        wret = tuple(map(lambda x:float(np.mean(x)),(t_drive, t_wait, prices)))
        # 返回前往最小权重充电站的路径和权重
        return stages[np.argmin(weight)].edges, wret # type: ignore
    
    def __start_trip(self, veh_id:str)->tuple[bool,Optional[TWeights]]:
        '''
        开始某辆车的当前行程
            veh_id: 车辆ID
        返回：
            发车成功: True, 快充站选择权重(如果需要途中快充)
            发车失败: False, None
        '''
        weights=None
        veh=self._VEHs[veh_id]
        trip=veh.trip
        direct_depart=True
        
        if ENABLE_DIST_BASED_CHARGING_DECISION:
            stage=self.__find_route(trip.depart_edge,trip.arrive_edge)
            # 判断电池电量是否充足
            direct_depart=veh.is_batt_enough(stage.length)
        else:
            # 判断电池电量是否需要快充
            direct_depart=veh.SOC>=veh.kfc
            if direct_depart: stage=self.__find_route(trip.depart_edge,trip.arrive_edge)
        if direct_depart: # 直接发车
            veh.target_CS=None
            veh.charge_target=veh.full_battery
            self.__add_veh(veh_id,stage.edges)
        else: # 中途充一次电
            route,weights=self.__sel_best_CS(veh,veh.omega,trip.depart_edge)
            if len(route)==0: # 电量不够开到任何一个充电站, 需要再充一段时间
                return False, None
            else: # 找到了充电站
                veh.target_CS=route[-1]
                self.__add_veh(veh_id,route)
        # 停止车辆的路边慢充并加入等待发车集合
        self._pk.pop_veh(veh_id)
        veh.stop_charging()
        veh.status=VehStatus.pending
        self.__pendSet.add(veh_id)
        return True,weights
    
    def __end_trip(self, veh_id:str):
        '''
        结束某辆车的当前行程, 并将其下次行程加入发车队列。如果行程终点具备充电条件, 则试图充电。
            veh: 车辆实例
        '''
        veh=self._VEHs[veh_id]
        veh.status=VehStatus.parking
        ctrip=veh.trip
        to_print=f"[{self.__ctime}, 到达]{veh.brief()}到达{ctrip.arrive_edge}, "
        if veh.SOC<veh.ksc:
            # 添加到路边慢充
            to_print+="加入慢充队列, " if self.__start_charging_SCS(veh) else "试图充电但无空位, "
        else:
            to_print+="不充电, "
        tid=veh.next_trip()
        if tid!=-1:
            ntrip=veh.trip
            self._que.push(ntrip.depart_time,veh_id)
            to_print+=f"下一行程{ntrip}."
        else:
            to_print+=f"所有行程结束."
        print(to_print,file=self.__ostream)
    
    def __start_charging_SCS(self, veh:EV) -> bool:
        '''
        使某辆车转入充电状态(路边慢充)
            veh: 车辆实例
        '''
        try:
            self._pk.add_veh(veh.ID,veh.trip.arrive_edge)
            veh.start_charging(False)
            return True
        except:
            return False
    
    def __start_charging_FCS(self, veh:EV):
        '''
        使某辆车转入充电状态(快充站)
            veh: 车辆实例
        '''
        veh.status=VehStatus.charging
        assert isinstance(veh.target_CS,str)
        if ENABLE_DIST_BASED_CHARGING_QUANTITY:
            ch_tar=veh.consumption*self.__find_route_static(veh.target_CS,veh.trip.arrive_edge).length
            if ch_tar>veh.full_battery: 
                # 即使再途中充满电, 也不足以到达终点
                print(f"[{self.__ctime}, 警告]{veh.brief()}的电池容量不足, 仅通过一次充电无法到达终点, 需要拖车.",file=self.__ostream)
            veh.charge_target=min(veh.full_battery,max(veh.full_battery*0.8,veh.krel*ch_tar))
        else:
            veh.charge_target=veh.full_battery
        veh.start_charging(True)
        self._cs.add_veh(veh.ID, veh.target_CS)                
        print(f"[{self.__ctime}, 充电]{veh.brief()}到达快充站{veh.target_CS}",file=self.__ostream)

    def __end_charging_FCS(self, veh:EV):
        '''
        使某辆车结束充电并发车(快充站)
            veh: 车辆实例
        '''
        if veh.target_CS is None:
            raise RuntimeError(f"运行时错误: {self.__ctime}, {veh.brief()}, {veh.status}")
        trip=veh.trip
        print(f"[{self.__ctime}, 充完]{veh.brief()}离开快充站{veh.target_CS}并前往{trip.arrive_edge}",file=self.__ostream)
        self.__add_veh(veh.ID,self.__find_route_static(veh.target_CS+"rev",trip.arrive_edge).edges)
        veh.target_CS=None
        veh.charge_target=veh.full_battery
        veh.status=VehStatus.pending
        veh.stop_charging()
        self.__pendSet.add(veh.ID)

    def __get_nearest_CS(self,cur_edge:str)->tuple[Optional[str],float,Optional[Stage]]:
        '''
        查找距离最近的充电站
            cur_edge: 当前所在道路
        '''
        min_cs_name=None
        min_cs_dist=1e400
        min_cs_stage=None
        for cs in self._cs.get_online_CS_names(self.__ctime):
            route=self.__find_route(cur_edge,cs)
            if route.length<min_cs_dist:
                min_cs_dist=route.length
                min_cs_name=cs
                min_cs_stage=route
        return min_cs_name,min_cs_dist,min_cs_stage

    def __batch_depart(self)->dict[str,Optional[TWeights]]:
        '''
        将发车队列中所有到时间的车辆全部发出
            self.__ctime: 当前时间, 单位秒
        返回：
            发车字典, 键为车辆ID, 值为快充站选择参数(如果无需前往快充站则是None)
        '''
        ret={}
        while not self._que.empty() and self._que.top[0]<=self.__ctime:
            depart_time,veh_id=self._que.pop()
            veh=self._VEHs[veh_id]
            trip=veh.trip
            success,weights=self.__start_trip(veh_id)
            if success:
                to_print=f"[{self.__ctime}, 出发]{veh.brief()}开始行程{veh.trip_id}: {trip.depart_edge}=>{trip.arrive_edge}."
                if weights is not None: to_print+=f" 将在{veh.target_CS}充电, Params={weights}."
                if depart_time<self.__ctime: to_print+=f" 延误{self.__ctime-depart_time}秒."
                print(to_print,file=self.__ostream)
                ret[veh_id]=weights
            else:
                cs_name,cs_dist,cs_stage=self.__get_nearest_CS(trip.depart_edge)
                batt_req=cs_dist*veh.consumption*veh.krel
                if self._pk.has_veh(veh.ID): # 插上了充电桩，可以等一等
                    delay=int(1+(batt_req-veh.battery)/veh.rate)
                    print(f"[{self.__ctime}, 延误]{veh_id}电量不足(当前:{veh.battery:.2f}, 需要:{batt_req:.2f}), 行程{veh.trip_id}发车失败, 等待{delay}秒再尝试",file=self.__ostream)
                    self._que.push(depart_time+delay,veh_id)
                else: #没插上了充电桩，传送到最近的快充站(消耗2倍运行时间)
                    veh.status = VehStatus.depleted
                    assert cs_name is not None and cs_stage is not None
                    veh.target_CS=cs_name
                    trT=int(self.__ctime+2*cs_stage.travelTime)
                    self._fQ.push(trT,veh.ID)
                    print(f"[{self.__ctime}, 故障]{veh_id}电量不足(当前:{veh.battery:.2f}, 需要:{batt_req:.2f}), 行程{veh.trip_id}发车失败, 将在{trT}传送至快充站{cs_name}进行充电",file=self.__ostream)
        return ret
    
    def __CS_update(self, sec:int):
        '''
        充电站更新: 对所有在充电站中的车辆进行充电, 并将充电完成的车辆发车
            sec: 仿真秒数
        '''
        veh_ids = self._cs.update(sec,self.__ctime)
        veh_ids.sort()
        for i in veh_ids:
            self.__end_charging_FCS(self._VEHs[i])
    
    def __PK_update(self, sec:int):
        '''
        停放车辆更新: 对所有在充电站中的停放车辆进行充电和V2G
            sec: 仿真秒数
        '''        
        self._pk.update(sec,self.__ctime)
    
    def get_sta_head(self)->list[str]:
        '''
        获取get_veh_count和CS_PK_update返回值对应的边名称
        '''
        return self.__names_cs+self.__names_ncs
    
    def get_veh_count(self)->list[int]:
        '''
        获取所有充电站和非充电站边中停放的车辆数
        '''
        return self._cs.get_veh_count()+self._pk.get_veh_count()
    
    def simulation_start(self, sumocfg_file:str, net_file:str, start_time:Optional[int]=None, gui:bool=True):
        '''
        开始仿真
            sumocfg_file: SUMO配置文件路径
            start_time: 开始时间(秒), 不指定则由SUMO配置文件提供
            gui: 是否显示图形化界面
        '''
        self.__gui=gui
        sumoCmd = ["sumo-gui" if self.__gui else "sumo", "-c", sumocfg_file, "-n", net_file, "--no-warnings", "--routing-algorithm","CH"]
        if start_time is not None: sumoCmd.extend(["-b", str(start_time)])
        traci.start(sumoCmd)

        self.__ctime = int(traci.simulation.getTime())
        self.__batch_depart()
    
    def simulation_step(self,step_len:int):
        '''
        仿真一步
            step_len: 仿真步长
            v2g_demand: V2G需求功率列表(kWh/s)
        '''
        traci.simulationStep(self.__ctime + step_len)
        new_time = int(traci.simulation.getTime())
        deltaT = new_time - self.__ctime
        self.__ctime = new_time
        
        cur_vehs:list[str]=traci.vehicle.getIDList()
        cur_vehs_set:set[str]=set(cur_vehs)

        # 将在SUMO中但状态为pending的车设为driving
        to_remove=[]
        for v in self.__pendSet:
            if v in cur_vehs_set:
                self._VEHs[v].status=VehStatus.driving
                self.__driveSet.add(v)
                to_remove.append(v)
        self.__pendSet.difference_update(to_remove)
        
        # 将不在SUMO中但状态为driving的车设为本行程结束/充电
        to_remove=[]
        for v in self.__driveSet:
            if not v in cur_vehs_set:
                to_remove.append(v)
        to_remove.sort()
        for v in to_remove:
            self.__driveSet.remove(v)
            if self._VEHs[v].target_CS==None:
                self.__end_trip(v)
            else:
                self.__start_charging_FCS(self._VEHs[v])
            
        # 处理在道路上行驶的车辆
        for veh_id in cur_vehs:
            veh=self._VEHs[veh_id]
            veh.drive(traci.vehicle.getDistance(veh_id))
            if veh.battery<=0: # 电量耗尽的车辆将被送往最近的快充站(花费时间*2)
                veh.status=VehStatus.depleted
                cur_edge = traci.vehicle.getRoadID(veh_id)
                veh.target_CS,_,cs_stage=self.__get_nearest_CS(cur_edge)
                assert cs_stage is not None and veh.target_CS is not None
                trT=int(self.__ctime+2*cs_stage.travelTime)
                self._fQ.push(trT,veh_id)
                traci.vehicle.remove(veh_id)
                print(f"[{self.__ctime}, 故障]{veh_id}电量耗尽, 已移出路网, 将在{trT}传送至快充站{veh.target_CS}进行充电",file=self.__ostream)

        # 处理在充电站和停放的车辆
        self.__CS_update(deltaT)
        self.__PK_update(deltaT)
        self.__batch_depart()

        # 处理故障车辆
        while not self._fQ.empty() and self._fQ.top[0]<=self.__ctime:
            _,v=self._fQ.pop()
            self.__start_charging_FCS(self._VEHs[v])
    
    def simulation_stop(self):
        if self.__gui is None:
            raise ValueError("仿真尚未开始, 请调用simulation_start以开始仿真")
        traci.close()